#!/usr/bin/Rscript
# setwd("/home/zhang/BMK/test/cuffdiff")
setwd("F:/cuffdiff/")

# 从cuffdiff的结果，读取要画热图的数据
read.fpkm <- function(input){
  if(!require(openxlsx)){
    install.packages("openxlsx")
    library(openxlsx)
  }

  input <- read.xlsx(as.character(input))
  input <- input[,c(8, 9, 14)]
  
  for(i in 1:(ncol(input) - 1)){
    input[input[,i] == 0, i] <- 0.01
  }
  colnames(input)[1:2] <- c("Control", "Treat")
  return(input)
}

draw.fpkm.heatmap <- function(input, name){
  if(!require(gplots)){
    install.packages("gplots")
    library(gplots)
  }

  # 如果输入格式是data.drame就转化成matrix，比较不容易出错
  if(is.data.frame(input)){
    input <- as.matrix(input)
  }

  # input <- scale(input, center = T, scale = T)
  input <- log(input)

  tiff(as.character(name), units = "in", width = 7, height = 15, res = 600, compression = "lzw")
  heatmap.2(input, dendrogram = "none", labRow = F, na.rm = T,
            col = redgreen(75), trace = "none", scale = "column",
            srtCol = 0, adjCol = c(NA, 1))
  dev.off()
}

mrna <- read.fpkm("mrna_fpkm.xlsx")
draw.fpkm.heatmap(mrna[,1:2], "mrna_fpkm_heatmap.tiff")

lnc <- read.fpkm("lnc_fpkm.xlsx")
draw.fpkm.heatmap(lnc[,1:2], "lnc_fpkm_heatmap.tiff")

mrna <- mrna[mrna$significant == "yes", 1:2]
draw.fpkm.heatmap(mrna, "mrna_fpkm_diff_heatmap.tiff")

lnc <- lnc[lnc$significant == "yes", 1:2]
draw.fpkm.heatmap(lnc, "lnc_fpkm_diff_heatmap.tiff")

########################################################################################################
# 单独做microRNA的热图

read.micro <- function(input, diff = FALSE){
  if(!require(openxlsx)){
    install.packages("openxlsx")
    if(!require(openxlsx)){
      stop("openxlsx install failed")
    }
  }
  
  if(!require(edgeR)){
    source("https://www.bioconductor.org/biocLite.R")
    biocLite("edgeR")
    if(!require(edgeR)){
      stop("edgeR install failed")
    }
  }
  input <- read.xlsx(as.character(input))
  
  if(diff){
    input <- input[abs(input$logFC) >= 1 & input$FDR < 0.05,]
  }
  
  input <- input[,c("control", "treat")]
  
  # 用来将表达量为0的值替换为0.01
  change.zero <- function(input){
    input[input == 0] <- 0.01
    return(input)
  }
  
  input <- cpm(input)
  input <- apply(input, 2, change.zero)
  return(input)
}

test <- read.micro("microRNA.xlsx")
draw.fpkm.heatmap(test, "microRNA_cpm_heatmap.tiff")

test <- read.micro("microRNA.xlsx", diff = T)
draw.fpkm.heatmap(test, "microRNA_cpm_diff_heatmap.tiff")
